/*************************************************************************
    > File Name: tracker1D_api.c
    > Author: wanghuanyu
    > Mail: wanghuanyu@weye.tech
    > Created Time: [2024-03-07 01:19:13]
 ************************************************************************/
#include "tracker1D_api.h"

#include "app/kalman.h"
#include "core/math/sys_math.h"
#include "core/common/memory.h"

#define _MAX_TOR_   1e5f

typedef struct tracker1D_data_prop_t_tag
{
	we_f32 m_r;
	we_f32 m_v;
}tracker1D_data_prop_t;

typedef struct tracker1D_data_t_tag
{
	we_u16  m_id;
	struct tracker1D_data_flag_t
	{
		we_u16 b_use : 1;
		we_u16 b_valid : 1;
		we_u16 b_dead : 1;
		we_u16 b_new : 1;
		we_u16 b_found : 1;
		we_u16 b_lost : 1;
		we_u16 m_rsvd : 10;
	}m_flag;

	struct tracker1D_data_stat_t
	{
		we_u8 m_live_cnt;
		we_u8 m_lost_cnt;
		we_u8 m_clive_cnt;
		we_u8 m_clost_cnt;
	}m_stats;

	tracker1D_data_prop_t m_props;

	kalman_rv_context_t m_kalman;

}tracker1D_data_t;

typedef struct tracker1D_input_data_t_tag
{
	tracker1D_data_prop_t m_props;

	we_f32 m_min_diff;
	tracker1D_data_t* m_pair;
}tracker1D_input_data_t;

typedef struct tracker1D_context_t_tag
{
	we_u32 m_curr_id;
	we_f32 m_dr;

	we_u32 m_output_threshold;
	we_u32 m_dead_threshold;
	we_u32 m_valid_dead_threshold;
	we_u32 m_valid_threshold;
	we_f32 m_tor;

	we_u8 m_buffer_cnt;
	tracker1D_input_data_t m_buffer[_MAX_INPUT_];

	we_u32 m_curr_cnt;
	tracker1D_data_t  m_curr[_MAX_INPUT_ << 2];
	tracker1D_data_t* m_curr_ptr[_MAX_INPUT_ << 2];

	tracker1D_output_t m_output;
}tracker1D_context_t;

static tracker1D_context_t gctx;

_WE_API void tracker1D_init(const we_f32 dr)
{
	_mem_reset(&gctx, sizeof(gctx));

	gctx.m_dr = dr;

	gctx.m_output_threshold = 2;
	gctx.m_dead_threshold = 2;
	gctx.m_valid_dead_threshold = 3;
	gctx.m_valid_threshold = 3;
	gctx.m_tor = 3.f * dr;
}

static void _tracker1D_process_update(
	tracker1D_data_t* data, const tracker1D_input_data_t* distances, const we_u32 n)
{
	//data->m_flag.b_use : 1;
	data->m_flag.b_dead = 0;
	data->m_flag.b_new = 0;
	data->m_flag.b_found = 1;
	data->m_flag.b_lost = 0;

	data->m_stats.m_live_cnt++;
	data->m_stats.m_lost_cnt;
	data->m_stats.m_clost_cnt = 0;

	we_f32 Z = data->m_props.m_r;

	we_f32 min_diff = 1e5f;
	const tracker1D_input_data_t* sel_d = we_null;
	for (we_u32 i = 0; i < n; ++i)
	{
		const tracker1D_input_data_t* d = &distances[i];
		if (d->m_pair == data)
		{
			we_f32 diff = _mat_abs(data->m_props.m_r - d->m_props.m_r);
			if (min_diff > diff)
			{
				min_diff = diff;
				sel_d = d;
			}
		}
	}
	if (sel_d != we_null)
	{
		Z = sel_d->m_props.m_r;
		data->m_stats.m_clive_cnt++;
		if (data->m_stats.m_clive_cnt >= gctx.m_valid_threshold)
		{
			data->m_flag.b_valid = 1;
		}
	}

	we_f32 rv[2];
	kalman_rv_apply(rv, &data->m_kalman, Z);
	data->m_props.m_r = rv[0];
	data->m_props.m_v = rv[1];
}

static void _tracker1D_process_lost(tracker1D_data_t* data)
{
	//data->m_flag.b_use : 1;
	data->m_flag.b_dead = 0;
	data->m_flag.b_new = 0;
	data->m_flag.b_found = 0;
	data->m_flag.b_lost = 1;

	data->m_stats.m_live_cnt++;
	data->m_stats.m_lost_cnt++;
	data->m_stats.m_clive_cnt = 0;
	data->m_stats.m_clost_cnt++;

	if (data->m_flag.b_valid == 1)
	{
		if (data->m_stats.m_clost_cnt >= gctx.m_valid_dead_threshold)
		{
			data->m_flag.b_dead = 1;
		}
	}
	else if (data->m_stats.m_clost_cnt >= gctx.m_dead_threshold)
	{
		data->m_flag.b_dead = 1;
	}

	if (data->m_flag.b_dead == 0)
	{
		we_f32 rv[2];
		kalman_rv_apply(rv, &data->m_kalman, data->m_props.m_r);
		data->m_props.m_r = rv[0];
		data->m_props.m_v = rv[1];
	}
}

static tracker1D_data_t* _tracker1D_generate_data(tracker1D_context_t* ctx,
	const tracker1D_input_data_t* inp, const we_f32 dt)
{
	static we_f32 Q[] = { 1e-7f, 0.f, 0.f, 1e-7f };
	tracker1D_data_t* ret = we_null;
	for (we_u32 i = 0; i < (_MAX_INPUT_ << 2); ++i)
	{
		tracker1D_data_t* data = &ctx->m_curr[i];
		if (data->m_flag.b_use == 0)
		{
			data->m_flag.b_use = 1;
			data->m_flag.b_valid = 0;
			data->m_flag.b_dead = 0;
			data->m_flag.b_found = 0;
			data->m_flag.b_lost = 0;
			data->m_flag.b_new = 1;

			data->m_props.m_r = inp->m_props.m_r;

			data->m_stats.m_live_cnt = 1;
			data->m_stats.m_lost_cnt = 0;
			data->m_stats.m_clive_cnt = 1;
			data->m_stats.m_clost_cnt = 0;

			data->m_id = ctx->m_curr_id;
			ctx->m_curr_id++;

			kalman_rv_init(&data->m_kalman, data->m_props.m_r, 0.f, dt, Q, ctx->m_dr * ctx->m_dr);

			ret = data;
			ctx->m_curr_ptr[ctx->m_curr_cnt++] = data;
			break;
		}
	}
	return ret;
}

_WE_API we_u32 tracker1D_tracking(const we_f32* distances, const we_u32 n, const we_f32 dt)
{
	// 0.1 copy to buffer
	we_u32 cnt = ((n > _MAX_INPUT_) ? _MAX_INPUT_ : n);
	for (we_u32 i = 0; i < cnt; i++)
	{
		tracker1D_input_data_t* inp = &gctx.m_buffer[i];
		inp->m_props.m_r = distances[i];
		inp->m_pair = we_null;
		inp->m_min_diff = _MAX_TOR_;
	}
	gctx.m_buffer_cnt = cnt;

	// 0.2 prepare pointers
	cnt = 0;
	for (we_u32 i = 0; i < (_MAX_INPUT_ << 2); ++i)
	{
		tracker1D_data_t* data = &gctx.m_curr[i];
		if (data->m_flag.b_use == 0)
		{
			continue;
		}
		if (data->m_flag.b_dead == 1)
		{
			_mem_reset(data, sizeof(tracker1D_data_t));
			continue;
		}
		gctx.m_curr_ptr[cnt++] = data;
	}
	for (we_u32 i = cnt; i < (_MAX_INPUT_ << 2); ++i)
	{
		gctx.m_curr_ptr[i] = we_null;
	}
	for (we_u32 i = 0; i < cnt; ++i)
	{
		for (we_u32 j = i + 1; j < cnt; ++j)
		{
			if (gctx.m_curr_ptr[i]->m_stats.m_live_cnt < gctx.m_curr_ptr[j]->m_stats.m_live_cnt)
			{
				tracker1D_data_t* tmp = gctx.m_curr_ptr[j];
				gctx.m_curr_ptr[j] = gctx.m_curr_ptr[i];
				gctx.m_curr_ptr[i] = tmp;
			}
		}
	}
	gctx.m_curr_cnt = cnt;

	// 1. search current
	for (we_u32 i = 0; i < gctx.m_curr_cnt; ++i)
	{
		tracker1D_data_t* data = gctx.m_curr_ptr[i];

		we_bool found = we_false;
		for (we_u32 j = 0; j < gctx.m_buffer_cnt; ++j)
		{
			tracker1D_input_data_t* inp = &gctx.m_buffer[j];

			we_f32 diff = _mat_abs(data->m_props.m_r - inp->m_props.m_r);

			if (diff < gctx.m_tor)
			{
				if (diff < inp->m_min_diff)
				{
					inp->m_pair = data;
					inp->m_min_diff = diff;
				}
				found = we_true;
			}
		}
		if (found == we_true)
		{
			data->m_flag.b_found = 1;
		}
		else
		{
			data->m_flag.b_found = 0;
		}
	}

	// 2. update
	for (we_u32 i = 0; i < gctx.m_curr_cnt; ++i)
	{
		tracker1D_data_t* data = gctx.m_curr_ptr[i];

		data->m_props.m_r += data->m_props.m_v;
		if (data->m_flag.b_found == 1)
		{
			_tracker1D_process_update(data, gctx.m_buffer, gctx.m_buffer_cnt);
		}
		else
		{
			_tracker1D_process_lost(data);
		}
	}
	// 4. generate
	for (we_u32 i = 0; i < gctx.m_buffer_cnt; ++i)
	{
		tracker1D_input_data_t* inp = &gctx.m_buffer[i];
		if (inp->m_pair == we_null)
		{
			_tracker1D_generate_data(&gctx, inp, dt);
		}
	}
	// 4. output
	we_u32 output_cnt = 0;
	for (we_u32 i = 0; i < gctx.m_curr_cnt; ++i)
	{
		tracker1D_data_t* data = gctx.m_curr_ptr[i];
		if (data->m_flag.b_dead == 1) continue;
		if (data->m_flag.b_valid == 0) continue;

		if (data->m_stats.m_live_cnt >= gctx.m_output_threshold && data->m_stats.m_live_cnt > 1)
		{
			gctx.m_output.m_data[output_cnt].m_distance = data->m_props.m_r;
			gctx.m_output.m_data[output_cnt].m_velocity = data->m_props.m_v;
			gctx.m_output.m_data[output_cnt].m_id = data->m_id;
			output_cnt++;
		}
	}
	gctx.m_output.m_cnt = output_cnt;

	return output_cnt;
}

_WE_API tracker1D_output_t* tracker1D_get_output()
{
	return &gctx.m_output;
}
